import json
import os
import re
import sys
import platform
import socket
import psutil
from typing import Dict, Any, List, Optional
from pprint import pprint

class SystemInfoCollector:
    """用于自动化测试的系统信息收集工具"""

    def __init__(self, categories: Optional[List[str]] = None):
        """
        初始化系统信息收集器

        Args:
            categories: 需要收集的信息类别列表，默认收集所有类别
        """
        self.system = platform.system()
        if self.system != "Linux":
            raise OSError("此脚本仅支持Linux系统")

        # 支持的信息类别
        self.supported_categories = ["system", "cpu", "memory", "disk", "network"]

        # 初始化信息字典
        self.info = {
            "system": {},
            "cpu": {},
            "memory": {},
            "disk": {},
            "network": {},
        }

        # 过滤需要收集的类别
        if categories:
            invalid_cats = [cat for cat in categories if cat not in self.supported_categories]
            if invalid_cats:
                raise ValueError(f"无效的信息类别: {', '.join(invalid_cats)}")
            self.categories = categories
        else:
            self.categories = self.supported_categories

    def collect_all_info(self) -> Dict[str, Any]:
        """收集所有指定类别的系统信息"""
        for category in self.categories:
            try:
                collector_method = getattr(self, f"_collect_{category}_info")
                collector_method()
            except Exception as e:
                self.info[category]["error"] = str(e)

        # 只返回请求的类别
        return {category: self.info[category] for category in self.categories}

    def _collect_system_info(self) -> None:
        """收集基本系统信息"""
        self.info["system"]["hostname"] = socket.gethostname()
        self.info["system"]["kernel"] = platform.release()
        self.info["system"]["distribution"] = self._get_linux_distribution()

        try:
            with open("/proc/uptime", "r") as f:
                uptime_seconds = float(f.readline().split()[0])
                self.info["system"]["uptime_seconds"] = uptime_seconds
        except Exception as e:
            self.info["system"]["uptime_seconds"] = None
            self.info["system"]["error"] = str(e)

    def _collect_cpu_info(self) -> None:
        """收集CPU信息"""
        self.info["cpu"]["architecture"] = platform.machine()

        try:
            with open("/proc/cpuinfo", "r") as f:
                cpuinfo = f.read()

            vendor = re.search(r"vendor_id\s+:\s+(.+)", cpuinfo)
            model = re.search(r"model name\s+:\s+(.+)", cpuinfo)
            self.info["cpu"]["vendor"] = vendor.group(1) if vendor else "Unknown"
            self.info["cpu"]["model"] = model.group(1) if model else "Unknown"

            self.info["cpu"]["physical_cores"] = len(re.findall(r"processor\s+:\s+\d+", cpuinfo))

            siblings = re.search(r"siblings\s+:\s+(\d+)", cpuinfo)
            self.info["cpu"]["logical_cores"] = int(siblings.group(1)) if siblings else 0

            freq = re.search(r"cpu MHz\s+:\s+([\d.]+)", cpuinfo)
            self.info["cpu"]["frequency_mhz"] = float(freq.group(1)) if freq else None

            self.info["cpu"]["usage_percent"] = psutil.cpu_percent(interval=0.5)

        except Exception as e:
            self.info["cpu"]["error"] = str(e)

    def _collect_memory_info(self) -> None:
        """收集内存信息"""
        try:
            mem = psutil.virtual_memory()
            swap = psutil.swap_memory()

            self.info["memory"]["total_bytes"] = mem.total
            self.info["memory"]["available_bytes"] = mem.available
            self.info["memory"]["used_percent"] = mem.percent

            self.info["memory"]["swap_total_bytes"] = swap.total
            self.info["memory"]["swap_used_bytes"] = swap.used

            with open("/proc/meminfo", "r") as f:
                meminfo = f.read()

            cached = re.search(r"Cached:\s+(\d+)", meminfo)
            buffers = re.search(r"Buffers:\s+(\d+)", meminfo)
            self.info["memory"]["cached_bytes"] = int(cached.group(1)) * 1024 if cached else None
            self.info["memory"]["buffers_bytes"] = int(buffers.group(1)) * 1024 if buffers else None

        except Exception as e:
            self.info["memory"]["error"] = str(e)

    def _collect_disk_info(self) -> None:
        """收集磁盘信息"""
        try:
            partitions = []
            for partition in psutil.disk_partitions():
                if os.name == 'nt':  # 跳过CD-ROM驱动器
                    if 'cdrom' in partition.opts or partition.fstype == '':
                        continue
                try:
                    usage = psutil.disk_usage(partition.mountpoint)
                    partition_info = {
                        "device": partition.device,
                        "mountpoint": partition.mountpoint,
                        "fstype": partition.fstype,
                        "total_bytes": usage.total,
                        "used_bytes": usage.used,
                        "free_bytes": usage.free,
                        "used_percent": usage.percent
                    }
                    partitions.append(partition_info)
                except PermissionError:
                    continue

            self.info["disk"]["partitions"] = partitions

            io_counters = psutil.disk_io_counters()
            self.info["disk"]["io_read_count"] = io_counters.read_count
            self.info["disk"]["io_write_count"] = io_counters.write_count
            self.info["disk"]["io_read_bytes"] = io_counters.read_bytes
            self.info["disk"]["io_write_bytes"] = io_counters.write_bytes

        except Exception as e:
            self.info["disk"]["error"] = str(e)

    def _collect_network_info(self) -> None:
        """收集网络信息"""
        try:
            try:
                s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                s.connect(("8.8.8.8", 80))
                self.info["network"]["ip_address"] = s.getsockname()[0]
                s.close()
            except Exception:
                self.info["network"]["ip_address"] = "Unknown"

            interfaces = {}
            for interface, addrs in psutil.net_if_addrs().items():
                interface_info = {
                    "addresses": [],
                    "mac": "Unknown"
                }
                for addr in addrs:
                    if addr.family == socket.AF_INET:
                        interface_info["addresses"].append({
                            "family": "IPv4",
                            "address": addr.address,
                            "netmask": addr.netmask
                        })
                    elif addr.family == socket.AF_INET6:
                        interface_info["addresses"].append({
                            "family": "IPv6",
                            "address": addr.address,
                            "netmask": addr.netmask
                        })
                    elif addr.family == psutil.AF_LINK:
                        interface_info["mac"] = addr.address
                interfaces[interface] = interface_info

            self.info["network"]["interfaces"] = interfaces

            connections = psutil.net_connections(kind='inet')
            self.info["network"]["tcp_connections"] = len([c for c in connections if c.type == socket.SOCK_STREAM])
            self.info["network"]["udp_connections"] = len([c for c in connections if c.type == socket.SOCK_DGRAM])

        except Exception as e:
            self.info["network"]["error"] = str(e)

    def _get_linux_distribution(self) -> str:
        """获取Linux发行版信息"""
        try:
            if os.path.exists("/etc/os-release"):
                with open("/etc/os-release", "r") as f:
                    for line in f:
                        if line.startswith("PRETTY_NAME="):
                            return line.split("=")[1].strip().strip('"')
            elif os.path.exists("/etc/lsb-release"):
                with open("/etc/lsb-release", "r") as f:
                    for line in f:
                        if line.startswith("DISTRIB_DESCRIPTION="):
                            return line.split("=")[1].strip().strip('"')
            return "Unknown"
        except:
            return "Unknown"

def get_system_info(categories: Optional[List[str]] = None) -> Dict[str, Any]:
    """
    获取系统信息的便捷函数

    Args:
        categories: 需要收集的信息类别列表，默认收集所有类别

    Returns:
        包含系统信息的字典
    """
    try:
        collector = SystemInfoCollector(categories)
        return collector.collect_all_info()
    except Exception as e:
        return {"error": str(e)}

# 示例用法
if __name__ == "__main__":
    # 收集所有信息
    from common.common_func import pretty_print
    all_info = get_system_info()
    pretty_print(all_info)

    # 仅收集CPU和内存信息
    cpu_mem_info = get_system_info(["cpu", "memory"])
    print("\nCPU和内存信息:")
    pretty_print(cpu_mem_info)
